import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;

public class aviaTest
{
	private static final double DIFF_THRESHOLD = 0.4;
	private static final double HIGH_VALUE_THRESHOLD = 0.96;
	private static final double LOW_VALUE_THRESHOLD = 0.04;
	private static final double SIG_LOWER_THRESHOLD = 0.001;

	private static HashMap<Integer, SNPData> m_snps = new HashMap<Integer, SNPData>();
	private static List<Integer> m_significant = new ArrayList<Integer>();
	private static List<PersonData> m_testCases = new ArrayList<PersonData>();

	private static int SNPSlength;

	public static void main(String[] args)
	{
		try
		{
			String positionFile = /*
								 * "C:\\Users\\Maoz\\workspace\\HG2011c\\trunk\\files\\"
								 * +
								 */"src\\chr1.pos"; // args[1];
			String YRIFile = /*
							 * "C:\\Users\\Maoz\\workspace\\HG2011c\\trunk\\files\\"
							 * +
							 */"src\\YRI_freqs_104.txt"; // args[2];
			String CEUFile = /*
							 * "C:\\Users\\Maoz\\workspace\\HG2011c\\trunk\\files\\"
							 * +
							 */"src\\CEU_freqs_106.txt"; // args[3];
			String testFile = /*
							 * "C:\\Users\\Maoz\\workspace\\HG2011c\\trunk\\files\\"
							 * +
							 */"src\\train_genos.txt"; // args[3];
			String testFile2 = /*
								 * "C:\\Users\\Maoz\\workspace\\HG2011c\\trunk\\files\\"
								 * +
								 */"src\\train_pops.txt"; // args[3];

			readSNPsData(positionFile, YRIFile, CEUFile);
			readTestData(testFile, testFile2);
			
			getTestDataStats();
			
			double avg = 0;
			int count = 0;
			
			for (int i = 0; i < m_testCases.size(); i++)
			//for (int i = 0; i < 12; i++)
			{
				if ( i % 3 != 0)
				{
					double personGuess = testPerson(i);
					avg += personGuess;
					System.out.println("Correctly guessed in person " + i + " is " + personGuess);
					count++;
				}
			}
			
			double totalAvg = avg / count;
			
			System.out.println("Total average over all cases is " + totalAvg);

			// debug printouts
			System.out.println("Total snps are " + m_snps.size());
			System.out.println("Total persons are " + m_testCases.size());

		} catch (Exception e)
		{
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}

	private static double testPerson(int personIndex) throws IOException
	{
		double[][] chancesMatrix = new double[3][SNPSlength];

		byte[] SNPArray = m_testCases.get(personIndex).getSNPs();
		byte[] SNPGuessesArray = m_testCases.get(personIndex).getResultSNPs();
		byte[] currentGuesses = new byte[SNPSlength];

		for (int i = 0; i < SNPSlength; i++)
		{
			SNPData data = m_snps.get(i);
			double pCEU = data.getCEU();
			double pYRI = data.getYRI();
			int SNPNumAtPosition = SNPArray[i];
			if (SNPNumAtPosition == 0)
			{
				chancesMatrix[0][i] = (1 - pYRI) * (1 - pYRI);
				chancesMatrix[1][i] = (1 - pYRI) * (1 - pCEU);
				chancesMatrix[2][i] = (1 - pCEU) * (1 - pCEU);
			} else if (SNPNumAtPosition == 1)
			{
				chancesMatrix[0][i] = pYRI * (1 - pYRI);
				chancesMatrix[1][i] = pYRI * (1 - pCEU) + pCEU * (1 - pYRI);
				chancesMatrix[2][i] = pCEU * (1 - pCEU);
			} else if (SNPNumAtPosition == 2)
			{
				chancesMatrix[0][i] = pYRI * pYRI;
				chancesMatrix[1][i] = pYRI * pCEU;
				chancesMatrix[2][i] = pCEU * pCEU;
			}
		}
		
		Set<Integer> possibilitiesSet = new HashSet<Integer>();
		possibilitiesSet.add(0);
		possibilitiesSet.add(1);
		possibilitiesSet.add(2);
		int currentGuess = -1;
		int previousGuess = -1;
		int previousGuessStartIndex = -1;
		int oneLastZeroIndex = 0;
		int twoLaztZeroIndex = 0;
		int zeroLastZeroIndex = 0;
		int currentGuessZeroIndex = 0;;
		boolean guessFound = false;
		
		for (int i = 0; i < SNPSlength; i++)
		{
			//we've reached last line of data
			if (i == SNPSlength - 1)
			{
				for (int j = previousGuessStartIndex; j < SNPSlength; j++)
				{
					currentGuesses[j] = (byte)previousGuess;
				}	
				continue;
			}
			if (chancesMatrix[0][i] < SIG_LOWER_THRESHOLD && chancesMatrix[1][i] < SIG_LOWER_THRESHOLD && chancesMatrix[2][i] < SIG_LOWER_THRESHOLD)
			{
				continue;
			}
			
			if (chancesMatrix[0][i] < SIG_LOWER_THRESHOLD)
			{
				zeroLastZeroIndex = i;
				possibilitiesSet.remove(0);
				if (possibilitiesSet.size() == 1)
				{
					currentGuess = possibilitiesSet.iterator().next();
					guessFound = true;
				}
			}
			if (chancesMatrix[1][i] < SIG_LOWER_THRESHOLD)
			{
				oneLastZeroIndex = i;
				possibilitiesSet.remove(1);
				if (possibilitiesSet.size() == 1)
				{
					currentGuess = possibilitiesSet.iterator().next();
					guessFound = true;
				}
			}
			if (chancesMatrix[2][i] < SIG_LOWER_THRESHOLD)
			{
				twoLaztZeroIndex = i;
				possibilitiesSet.remove(2);
				if (possibilitiesSet.size() == 1)
				{
					currentGuess = possibilitiesSet.iterator().next();
					guessFound = true;
				}
			}
			
			// all lines has zeroes, skip
			if (possibilitiesSet.size() == 0)
			{
				guessFound = false;
				possibilitiesSet.add(0);
				possibilitiesSet.add(1);
				possibilitiesSet.add(2);
				continue;
			}
			
			if (guessFound)
			{
				if (currentGuess == previousGuess)
				{
					guessFound = false;
					possibilitiesSet.add(0);
					possibilitiesSet.add(1);
					possibilitiesSet.add(2);
					continue;
				}
				
				if (currentGuess == 0)
				{
					currentGuessZeroIndex = zeroLastZeroIndex;
				}
				else if (currentGuess == 1)
				{
					currentGuessZeroIndex = oneLastZeroIndex;
				}
				else
				{
					currentGuessZeroIndex = twoLaztZeroIndex;
				}
				
				if (previousGuess != -1)
				{	
					int maxProbsIndex = -1;
					double maxProbs = 0;
					
					for (int j = currentGuessZeroIndex + 1; j < i; j++)
					{
						double probabsProduct = 1.0;
						
						for (int h = currentGuessZeroIndex + 1; h < j; h++)
						{
							if (chancesMatrix[previousGuess][h] < SIG_LOWER_THRESHOLD)
							{
								continue;
							}
							probabsProduct *= chancesMatrix[previousGuess][h];
						}
						for (int h = j; h < i; h++)
						{
							if (chancesMatrix[previousGuess][h] < SIG_LOWER_THRESHOLD)
							{
								continue;
							}
							probabsProduct *= chancesMatrix[currentGuess][h];
						}
						
						if (probabsProduct > maxProbs)
						{
							maxProbs = probabsProduct;
							maxProbsIndex = j;
							probabsProduct = 1;
						}
					}
					
					if (maxProbsIndex == -1)
					{
						maxProbsIndex = i;
					}
					
					// fill in values in guess arrays
					for (int j = previousGuessStartIndex; j < maxProbsIndex; j++)
					{
						currentGuesses[j] = (byte)previousGuess;
					}
					previousGuessStartIndex = maxProbsIndex;
				}
				else
				{
					previousGuessStartIndex = 0;
				}
				
				previousGuess = currentGuess;
				possibilitiesSet.add(0);
				possibilitiesSet.add(1);
				possibilitiesSet.add(2);
				guessFound = false;
			}
		}
		
		int errorCount = 0;
		for (int i = 0; i < SNPSlength; i++)
		{
			if (currentGuesses[i] != SNPGuessesArray[i])
			{
				errorCount++;
			}
		}
		
		
		int correctGuesses = SNPSlength - errorCount;
		double totalSuccess = (double)correctGuesses / SNPSlength;

		if (personIndex == 0)
		{
			printFile("PersonOneGuesses.txt", currentGuesses);
		}
		
		return totalSuccess;
	}
	
	private static void getTestDataStats()
	{
		int maxLength = 0;
		int minLength = Integer.MAX_VALUE;
		int counter = 0;
		
		for (int i = 0; i < m_testCases.size(); i++)
		{
			byte[] SNPResultsArray = m_testCases.get(i).getResultSNPs();
			counter = 0;
			
			for (int j = 0; j < SNPSlength; j++)
			{
				if (j != 0)
				{
					if (SNPResultsArray[j] != SNPResultsArray[j-1])
					{
						if (counter > maxLength)
						{
							maxLength = counter;
						}
						if (counter < minLength)
						{
							minLength = counter;
						}
						counter = 0;
					}
				}
				counter++;	
			}
		}
		
		System.out.println("Shortest strand count is " + minLength);
		System.out.println("Longest strand count is " + maxLength);
	}

	private static void readTestData(String testFile, String testFile2)
			throws FileNotFoundException, IOException
	{
		int counter = 0;

		FileReader fr = new FileReader(testFile);

		BufferedReader br = new BufferedReader(fr);
		String s;
		while ((s = br.readLine()) != null)
		{
			s = s.trim();
			String[] snps = s.split("  ");
			byte[] data = new byte[snps.length];

			int i = 0;
			for (String snp : snps)
			{
				data[i] = Byte.parseByte(snp);
				i++;
			}

			PersonData person = new PersonData(counter);
			person.setSNPs(data);
			m_testCases.add(person);
			counter++;
		}
		fr.close();

		counter = 0;

		fr = new FileReader(testFile2);

		br = new BufferedReader(fr);
		while ((s = br.readLine()) != null)
		{
			s = s.trim();
			String[] snps = s.split("  ");
			byte[] data = new byte[snps.length];

			int i = 0;
			for (String snp : snps)
			{
				data[i] = Byte.parseByte(snp);
				i++;
			}

			PersonData person = m_testCases.get(counter);
			person.setResultSNPs(data);
			m_testCases.add(person);
			counter++;
		}
		fr.close();

	}

	private static void readSNPsData(String positionFile, String YRIFile,
			String CEUFile) throws FileNotFoundException, IOException
	{

		FileReader fr;
		int counter = 0;

		// get positions
		fr = new FileReader(positionFile);

		BufferedReader br = new BufferedReader(fr);
		String s;
		while ((s = br.readLine()) != null)
		{
			SNPData data = new SNPData(counter);
			data.setPosition(Integer.parseInt(s.trim()));
			m_snps.put(counter, data);
			counter++;
		}
		fr.close();
		SNPSlength = counter;

		// get CEU data
		counter = 0;
		fr = new FileReader(CEUFile);
		br = new BufferedReader(fr);
		while ((s = br.readLine()) != null)
		{
			m_snps.get(counter).setCEU(Float.parseFloat(s.trim()));
			counter++;
		}
		fr.close();

		// get YRI data
		counter = 0;
		fr = new FileReader(YRIFile);
		br = new BufferedReader(fr);
		while ((s = br.readLine()) != null)
		{
			m_snps.get(counter).setYRI(Float.parseFloat(s.trim()));
			counter++;
		}
		fr.close();
	}
	
	private static void printFile(String filename, byte[] arrayToPrint) throws IOException
	{
		int counter = 0;

		FileWriter fw = new FileWriter(filename);

		BufferedWriter bw = new BufferedWriter(fw);
		String s;
		
		int x = arrayToPrint.length;
		
		for (int i = 0; i < arrayToPrint.length ; i++)
		{
			int toWrite = arrayToPrint[i];
			bw.write(toWrite + "");
			bw.newLine();
		}
		
		fw.close();
	}
}
